import math
import random
import pygame
import rospy

import rvo.rvo_math as rvo_math
from rvo.vector import Vector2
from rvo.simulator import Simulator

RVO_RENDER = False


class Circle:
    def __init__(self):
        self.goals_ = []  # Vector2
        self.obstacles_ = []  # Vector2
        self.simulator_ = Simulator()

    def setup_scenario(self):
        # Specify the global time step of the simulation.
        self.simulator_.set_time_step(0.2)

        # Specify the default parameters for agents that are subsequently added.
        self.simulator_.set_agent_defaults(30.0, 10, 10.0, 10.0, 1.5, 2.0, Vector2(0.0, 0.0))

        # Add agents, specifying their start position, and store their goals on the opposite side of the environment.
        agent_num = 250
        for i in range(agent_num):
            self.simulator_.add_agent(200.0 *
                                      Vector2(math.cos(i * 2.0 * math.pi / agent_num),
                                              math.sin(i * 2.0 * math.pi / agent_num)),
                                      -200.0 *
                                      Vector2(math.cos(i * 2.0 * math.pi / agent_num),
                                              math.sin(i * 2.0 * math.pi / agent_num))
                                      )
            self.goals_.append(-self.simulator_.agents_[i].position_)

        # Add (polygonal) obstacles, specifying their vertices in counterclockwise order.
        obstacle1 = [Vector2(-10.0, 40.0), Vector2(-40.0, 40.0), Vector2(-40.0, 10.0), Vector2(-10.0, 10.0)]
        self.simulator_.add_obstacle(obstacle1)
        self.obstacles_.append(obstacle1)

        obstacle2 = [Vector2(10.0, 40.0), Vector2(10.0, 10.0), Vector2(40.0, 10.0), Vector2(40.0, 40.0)]
        self.simulator_.add_obstacle(obstacle2)
        self.obstacles_.append(obstacle2)

        obstacle3 = [Vector2(10.0, -40.0), Vector2(40.0, -40.0), Vector2(40.0, -10.0), Vector2(10.0, -10.0)]
        self.simulator_.add_obstacle(obstacle3)
        self.obstacles_.append(obstacle3)

        obstacle4 = [Vector2(-10.0, -40.0), Vector2(-10.0, -10.0), Vector2(-40.0, -10.0), Vector2(-40.0, -40.0)]
        self.simulator_.add_obstacle(obstacle4)
        self.obstacles_.append(obstacle4)

        # Process the obstacles so that they are accounted for in the simulation.
        self.simulator_.process_obstacles()

    def transform_position(self, position, width, height, bounds):
        # Transform the position to fit within the window bounds.
        x_min, x_max, y_min, y_max = bounds

        x = (position.x - x_min) / (x_max - x_min) * width
        y = (position.y - y_min) / (y_max - y_min) * height

        return x, height - y  # Flip y-axis for correct display

    def update_visualization(self, screen, width, height, bounds):

        screen.fill((255, 255, 255))

        for i in range(self.simulator_.num_agents):
            position = self.simulator_.agents_[i].position_
            transformed_radius = (width * self.simulator_.default_agent_.radius_) / (bounds[1] - bounds[0])
            transformed_position = self.transform_position(position, width, height, bounds)
            color = [0, 0, 0]
            color[i % 3] = 255  # Set one of the color channels to max (255) based on the agent index
            pygame.draw.circle(screen, color, transformed_position, transformed_radius)

        for obstacle in self.obstacles_:
            points = [self.transform_position(vec, width, height, bounds) for vec in obstacle]
            pygame.draw.polygon(screen, (0, 0, 0), points, 0)  # Fill the polygon with black color

    def set_preferred_velocities(self):
        """
        Set the preferred velocity to be a vector of unit magnitude (speed) in the direction of the goal.
        """
        collision_num = 0
        collision_agent = []
        for i in range(self.simulator_.num_agents):
            goal_vector = self.goals_[i] - self.simulator_.agents_[i].position_
            distSq = rvo_math.abs_sq(goal_vector)
            # if distSq > 1.0:
            #     goal_vector = rvo_math.normalize(goal_vector)
            #
            # self.simulator_.set_agent_pref_velocity(i, goal_vector)

            # # Perturb a little to avoid deadlocks due to perfect symmetry.
            # angle = random.random() * 2.0 * math.pi
            # dist = random.random() * 0.0001
            #
            # self.simulator_.set_agent_pref_velocity(i, self.simulator_.agents_[i].pref_velocity_ + dist * Vector2(
            #     math.cos(angle), math.sin(angle)))
            # print('agent' + str(self.simulator_.agents_[i].id_) + ' 终点距离: ' + str(math.sqrt(distSq)))
            if self.simulator_.agents_[i].collision_:
                collision_num += 1
                collision_agent.append(i)
        # print("collided num: ", collision_num, "collided ids: ", collision_agent)

    def reached_goal(self):
        """
        Check if all agents have reached their goals.
        """
        for i in range(self.simulator_.num_agents):
            if rvo_math.abs_sq(self.simulator_.agents_[i].position_ - self.goals_[i]
                               ) > self.simulator_.agents_[i].dist_update_now_goal**2:
                return False

        return True


def main():
    if RVO_RENDER:
        pygame.init()
        screen = pygame.display.set_mode((750, 750), pygame.RESIZABLE)
        clock = pygame.time.Clock()

    blocks = Circle()

    # Set up the scenario.
    blocks.setup_scenario()

    # Perform (and manipulate) the simulation.
    running = True
    bounds = (-225, 225, -225, 225)  # x_min, x_max, y_min, y_max
    step = 0
    while running and not blocks.reached_goal() and not rospy.is_shutdown():
        if RVO_RENDER:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.VIDEORESIZE:
                    screen = pygame.display.set_mode((event.w, event.h), pygame.RESIZABLE)

            width, height = screen.get_size()
            blocks.update_visualization(screen, width, height, bounds)

            pygame.display.flip()
            clock.tick(60)  # 控制帧率为60帧每秒

        blocks.set_preferred_velocities()
        blocks.simulator_.step()

        step += 1
        print("current step: ", step)

    pygame.quit()


if __name__ == '__main__':
    main()
